'''
Author: duliang thinktanker@163.com
Date: 2025-08-24 14:01:43
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-08-24 22:21:48
FilePath: \control-net\organize_files.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
import os
import shutil
from pathlib import Path


def organize_txt_and_jpg_files(source_folder, classes_folder, images_folder):
    """
    将txt文件存入classes文件夹，与txt文件名相同的jpg文件存入images文件夹
    
    Args:
        source_folder (str): 源文件夹路径
        classes_folder (str): txt文件目标文件夹路径
        images_folder (str): jpg文件目标文件夹路径
    """
    # 创建目标文件夹（如果不存在）
    Path(classes_folder).mkdir(parents=True, exist_ok=True)
    Path(images_folder).mkdir(parents=True, exist_ok=True)

    # 遍历源文件夹中的所有文件
    for root, dirs, files in os.walk(source_folder):
        for file in files:
            file_path = os.path.join(root, file)

            # 处理txt文件
            if file.endswith('.txt') and file != 'classes.txt':
                destination_path = os.path.join(classes_folder, file)
                # # 如果目标文件存在则先删除再剪切
                # if os.path.exists(destination_path):
                #     os.remove(destination_path)
                shutil.move(file_path, destination_path)
                print(f"已剪切txt文件: {file}")

            # 处理jpg文件
            elif file.endswith('.jpg'):
                # 检查是否有同名的txt文件
                txt_file = file.replace('.jpg', '.txt')
                txt_path = os.path.join(source_folder, txt_file)

                # 如果存在同名txt文件，则剪切jpg文件
                if os.path.exists(txt_path):
                    destination_path = os.path.join(images_folder, file)
                    # # 如果目标文件存在则先删除再剪切
                    # if os.path.exists(destination_path):
                    #     os.remove(destination_path)
                    shutil.move(file_path, destination_path)
                    print(f"已剪切jpg文件: {file} (存在对应的txt文件)")


# 添加新函数：找出两个文件夹中不对应的文件名
def find_mismatched_files(label_dir, images_dir):
    """
    找出label_dir和images_dir中不对应的文件名
    
    Args:
        label_dir (str): 包含txt标签文件的文件夹路径
        images_dir (str): 包含jpg图像文件的文件夹路径
    """
    # 获取label_dir中的所有txt文件名（不含扩展名）
    label_files = set()
    for file in os.listdir(label_dir):
        if file.endswith('.txt'):
            label_files.add(os.path.splitext(file)[0])
    
    # 获取images_dir中的所有jpg文件名（不含扩展名）
    image_files = set()
    for file in os.listdir(images_dir):
        if file.endswith('.jpg'):
            image_files.add(os.path.splitext(file)[0])
    
    # 找出只在label_dir中存在的文件（没有对应图片的标签文件）
    only_in_labels = label_files - image_files
    # 找出只在images_dir中存在的文件（没有对应标签的图片文件）
    only_in_images = image_files - label_files
    
    print("以下标签文件没有对应的图片文件:")
    for file in only_in_labels:
        print(f"  {file}.txt")
    
    print("\n以下图片文件没有对应的标签文件:")
    for file in only_in_images:
        print(f"  {file}.jpg")
    
    return only_in_labels, only_in_images


if __name__ == "__main__":
    # # 使用示例
    source_dir = r"D:\yoloimages"
    label_dir = r"D:\00myolo\labels"
    images_dir = r"D:\00myolo\images"

    organize_txt_and_jpg_files(source_dir, label_dir, images_dir)
    
    # 调用新函数找出不对应的文件
    # find_mismatched_files(label_dir, images_dir)